#!/usr/bin/env node

/**
 * PowerShell MCP服务器
 * 
 * 这是一个兼容PowerShell和CMD环境的MCP（Model Context Protocol）服务器实现。
 * 该服务器通过Stdio协议进行通信，提供PowerShell命令执行、系统信息获取等功能。
 * 
 * 主要功能：
 * 1. 执行PowerShell命令并返回结果
 * 2. 获取系统信息
 * 3. 文件和目录操作
 * 4. 进程管理
 * 5. 环境变量操作
 * 
 * 兼容性：
 * - Windows PowerShell 5.x
 * - PowerShell Core 7.x+
 * - CMD命令提示符
 * - 跨平台支持（Windows、Linux、macOS）
 * 
 * 使用方法：
 * 1. 构建项目：npm run build
 * 2. 启动服务器：npm start
 * 3. 或者直接运行：node build/index.js
 * 
 * @author MCP开发团队
 * @version 1.0.0
 */

import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { z } from "zod";
import { spawn, exec } from "child_process";
import { promisify } from "util";
import os from "os";
import fs from "fs/promises";
import path from "path";

// 将exec转换为Promise版本
const execAsync = promisify(exec);

/**
 * 检测当前操作系统类型
 * @returns 操作系统类型字符串
 */
function getPlatform(): string {
    const platform = os.platform();
    switch (platform) {
        case 'win32': return 'Windows';
        case 'darwin': return 'macOS';
        case 'linux': return 'Linux';
        default: return platform;
    }
}

/**
 * 检测可用的PowerShell版本
 * @returns PowerShell版本信息
 */
async function detectPowerShellVersion(): Promise<string> {
    try {
        // 优先检测PowerShell Core
        const { stdout } = await execAsync('pwsh --version');
        return stdout.trim();
    } catch {
        try {
            // 检测Windows PowerShell
            const { stdout } = await execAsync('powershell -Command "$PSVersionTable.PSVersion.ToString()"');
            return `Windows PowerShell ${stdout.trim()}`;
        } catch {
            return 'CMD (Command Prompt)';
        }
    }
}

/**
 * 执行命令的通用函数，自动适配不同环境
 * @param command 要执行的命令
 * @param shellType 指定的shell类型（可选）
 * @returns 执行结果
 */
async function executeCommand(command: string, shellType?: string): Promise<{
    stdout: string;
    stderr: string;
    exitCode: number;
    shellUsed: string;
}> {
    let shell: string;
    let shellArgs: string[];

    // 根据shell类型或自动检测选择shell
    if (shellType) {
        shell = shellType.toLowerCase();
    } else {
        // 自动检测最佳shell
        try {
            await execAsync('pwsh --version');
            shell = 'pwsh';
        } catch {
            try {
                await execAsync('powershell -Command "Get-Host"');
                shell = 'powershell';
            } catch {
                shell = process.platform === 'win32' ? 'cmd' : 'bash';
            }
        }
    }

    // 设置shell参数
    switch (shell) {
        case 'pwsh':
            shell = 'pwsh';
            shellArgs = ['-NoLogo', '-NoProfile', '-Command', command];
            break;
        case 'powershell':
            shell = 'powershell';
            shellArgs = ['-NoLogo', '-NoProfile', '-Command', command];
            break;
        case 'cmd':
            shell = 'cmd';
            shellArgs = ['/c', command];
            break;
        case 'bash':
        default:
            shell = 'bash';
            shellArgs = ['-c', command];
            break;
    }

    try {
        const { stdout, stderr } = await execAsync(command, {
            shell: shell === 'cmd' ? 'cmd.exe' : shell === 'powershell' ? 'powershell.exe' : shell === 'pwsh' ? 'pwsh.exe' : shell
        });
        
        return {
            stdout: stdout || '',
            stderr: stderr || '',
            exitCode: 0,
            shellUsed: shell
        };
    } catch (error: any) {
        return {
            stdout: error.stdout || '',
            stderr: error.stderr || error.message || '执行失败',
            exitCode: error.code || 1,
            shellUsed: shell
        };
    }
}

/**
 * 获取系统信息的辅助函数
 * @returns 系统信息对象
 */
async function getSystemInfo(): Promise<{
    platform: string;
    hostname: string;
    username: string;
    homeDir: string;
    tempDir: string;
    uptime: number;
    memory: {
        total: number;
        free: number;
        used: number;
    };
    cpus: {
        model: string;
        speed: number;
        cores: number;
    }[];
    networkInterfaces: any;
}> {
    const cpus = os.cpus();
    const totalMemory = os.totalmem();
    const freeMemory = os.freemem();

    return {
        platform: getPlatform(),
        hostname: os.hostname(),
        username: os.userInfo().username,
        homeDir: os.homedir(),
        tempDir: os.tmpdir(),
        uptime: os.uptime(),
        memory: {
            total: totalMemory,
            free: freeMemory,
            used: totalMemory - freeMemory
        },
        cpus: cpus.map(cpu => ({
            model: cpu.model,
            speed: cpu.speed,
            cores: cpus.length
        })),
        networkInterfaces: os.networkInterfaces()
    };
}

/**
 * 格式化文件大小
 * @param bytes 字节数
 * @returns 格式化的文件大小字符串
 */
function formatFileSize(bytes: number): string {
    const sizes = ['B', 'KB', 'MB', 'GB', 'TB'];
    if (bytes === 0) return '0 B';
    const i = Math.floor(Math.log(bytes) / Math.log(1024));
    return Math.round(bytes / Math.pow(1024, i) * 100) / 100 + ' ' + sizes[i];
}

// 创建MCP服务器实例
const server = new McpServer({
    name: "powershell-mcp",
    version: "1.0.0",
    description: "兼容PowerShell和CMD的MCP服务器，提供系统管理和命令执行功能",
    capabilities: {
        tools: {},
        resources: {}
    }
});

/**
 * 注册PowerShell命令执行工具
 * 
 * 功能：执行任意PowerShell、CMD或其他shell命令
 * 参数：
 * - command: 要执行的命令字符串
 * - shellType: 可选的shell类型（pwsh/powershell/cmd/bash）
 * 返回值：命令执行结果，包括stdout、stderr、退出码和使用的shell类型
 */
server.tool(
    "execute-command",
    "执行任意命令，自动适配PowerShell、CMD或其他shell环境",
    {
        command: z.string().describe("要执行的命令字符串"),
        shellType: z.string().optional().describe("可选的shell类型：pwsh、powershell、cmd、bash")
    },
    async ({ command, shellType }) => {
        const result = await executeCommand(command, shellType);
        
        return {
            content: [
                {
                    type: "text",
                    text: `命令执行完成

使用的Shell: ${result.shellUsed}
退出码: ${result.exitCode}

标准输出:
${result.stdout || '(无输出)'}

标准错误:
${result.stderr || '(无错误)'}`
                }
            ]
        };
    }
);

/**
 * 注册获取系统信息工具
 * 
 * 功能：获取详细的系统信息，包括硬件、操作系统、网络等
 * 参数：无
 * 返回值：格式化的系统信息报告
 */
server.tool(
    "get-system-info",
    "获取详细的系统信息，包括操作系统、硬件配置、网络接口等",
    {},
    async () => {
        const systemInfo = await getSystemInfo();
        const powerShellVersion = await detectPowerShellVersion();
        
        return {
            content: [
                {
                    type: "text",
                    text: `系统信息报告

操作系统: ${systemInfo.platform}
主机名: ${systemInfo.hostname}
用户名: ${systemInfo.username}
PowerShell版本: ${powerShellVersion}

目录信息:
主目录: ${systemInfo.homeDir}
临时目录: ${systemInfo.tempDir}

系统运行时间: ${Math.floor(systemInfo.uptime / 3600)}小时${Math.floor((systemInfo.uptime % 3600) / 60)}分钟

内存信息:
总内存: ${formatFileSize(systemInfo.memory.total)}
可用内存: ${formatFileSize(systemInfo.memory.free)}
已用内存: ${formatFileSize(systemInfo.memory.used)}
使用率: ${Math.round((systemInfo.memory.used / systemInfo.memory.total) * 100)}%

CPU信息:
${systemInfo.cpus.map((cpu, index) => 
    `CPU ${index + 1}: ${cpu.model} @ ${cpu.speed} MHz (${cpu.cores}核)`
).join('\n')}

网络接口:
${Object.entries(systemInfo.networkInterfaces)
    .map(([name, interfaces]) => {
        const ips = (interfaces as any[])
            .filter(iface => !iface.internal && iface.family === 'IPv4')
            .map(iface => iface.address)
            .join(', ');
        return ips ? `${name}: ${ips}` : null;
    })
    .filter(Boolean)
    .join('\n') || '无网络接口信息'}`
                }
            ]
        };
    }
);

/**
 * 注册文件操作工具
 * 
 * 功能：执行文件和目录操作
 * 参数：
 * - operation: 操作类型（list/read/write/delete/info）
 * - path: 文件或目录路径
 * - content: 写入内容（仅write操作需要）
 * 返回值：操作结果
 */
server.tool(
    "file-operation",
    "执行文件和目录操作，支持列出、读取、写入、删除和获取文件信息",
    {
        operation: z.enum(["list", "read", "write", "delete", "info"]).describe("操作类型"),
        path: z.string().describe("文件或目录路径"),
        content: z.string().optional().describe("写入内容（仅write操作需要）")
    },
    async ({ operation, path: filePath, content }) => {
        try {
            switch (operation) {
                case "list":
                    const files = await fs.readdir(filePath, { withFileTypes: true });
                    const fileList = [];
                    
                    for (const item of files) {
                        const fileInfo = {
                            name: item.name,
                            type: item.isDirectory() ? "directory" : item.isFile() ? "file" : "other",
                            size: "-"
                        };
                        
                        if (item.isFile()) {
                            try {
                                const stat = await fs.stat(path.join(filePath, item.name));
                                fileInfo.size = formatFileSize(stat.size);
                            } catch {
                                fileInfo.size = "N/A";
                            }
                        }
                        
                        fileList.push(fileInfo);
                    }
                    
                    return {
                        content: [
                            {
                                type: "text",
                                text: `目录列表: ${filePath}

${fileList.map(f => `${f.type === 'directory' ? '📁' : '📄'} ${f.name} ${f.type === 'file' ? `(${f.size})` : ''}`).join('\n')}`
                            }
                        ]
                    };

                case "read":
                    const fileContent = await fs.readFile(filePath, 'utf-8');
                    return {
                        content: [
                            {
                                type: "text",
                                text: `文件内容: ${filePath}

${fileContent}`
                            }
                        ]
                    };

                case "write":
                    if (!content) {
                        throw new Error("写入操作需要提供content参数");
                    }
                    await fs.writeFile(filePath, content, 'utf-8');
                    return {
                        content: [
                            {
                                type: "text",
                                text: `文件写入成功: ${filePath}`
                            }
                        ]
                    };

                case "delete":
                    const stat = await fs.stat(filePath);
                    if (stat.isDirectory()) {
                        await fs.rm(filePath, { recursive: true, force: true });
                    } else {
                        await fs.unlink(filePath);
                    }
                    return {
                        content: [
                            {
                                type: "text",
                                text: `删除成功: ${filePath}`
                            }
                        ]
                    };

                case "info":
                    const fileStat = await fs.stat(filePath);
                    return {
                        content: [
                            {
                                type: "text",
                                text: `文件信息: ${filePath}

类型: ${fileStat.isDirectory() ? '目录' : fileStat.isFile() ? '文件' : '其他'}
大小: ${formatFileSize(fileStat.size)}
创建时间: ${fileStat.birthtime.toLocaleString()}
修改时间: ${fileStat.mtime.toLocaleString()}
访问时间: ${fileStat.atime.toLocaleString()}
权限: ${fileStat.mode.toString(8)}`
                            }
                        ]
                    };
            }
        } catch (error: any) {
            return {
                content: [
                    {
                        type: "text",
                        text: `操作失败: ${error.message}`
                    }
                ]
            };
        }
    }
);

/**
 * 注册进程管理工具
 * 
 * 功能：获取运行中的进程信息
 * 参数：
 * - filter: 进程名过滤字符串（可选）
 * 返回值：进程列表信息
 */
server.tool(
    "get-processes",
    "获取当前运行的进程列表，支持按名称过滤",
    {
        filter: z.string().optional().describe("进程名称过滤字符串")
    },
    async ({ filter }) => {
        try {
            let command: string;
            
            if (process.platform === 'win32') {
                command = 'tasklist /FO CSV /NH';
                if (filter) {
                    command += ` | findstr /I "${filter}"`;
                }
            } else {
                command = 'ps aux';
                if (filter) {
                    command += ` | grep -i "${filter}"`;
                }
            }
            
            const result = await executeCommand(command);
            
            if (result.exitCode === 0) {
                const lines = result.stdout.trim().split('\n');
                const processes = lines.slice(1).map(line => {
                    if (process.platform === 'win32') {
                        const parts = line.split('","');
                        return {
                            name: parts[0]?.replace(/"/g, '') || 'N/A',
                            pid: parts[1]?.replace(/"/g, '') || 'N/A',
                            memory: parts[4]?.replace(/"/g, '') || 'N/A'
                        };
                    } else {
                        const parts = line.trim().split(/\s+/);
                        return {
                            user: parts[0] || 'N/A',
                            pid: parts[1] || 'N/A',
                            cpu: parts[2] || 'N/A',
                            memory: parts[3] || 'N/A',
                            command: parts.slice(10).join(' ') || 'N/A'
                        };
                    }
                });
                
                return {
                    content: [
                        {
                            type: "text",
                            text: `进程列表${filter ? ` (过滤: ${filter})` : ''}

${processes.slice(0, 50).map(p => 
    process.platform === 'win32' 
        ? `${p.name} (PID: ${p.pid}) - 内存: ${p.memory}`
        : `${p.command} (PID: ${p.pid}, CPU: ${p.cpu}%, 内存: ${p.memory}%)`
).join('\n')}

${processes.length > 50 ? `... 共${processes.length}个进程，显示前50个` : `共${processes.length}个进程`}`
                        }
                    ]
                };
            } else {
                return {
                    content: [
                        {
                            type: "text",
                            text: `获取进程信息失败: ${result.stderr}`
                        }
                    ]
                };
            }
        } catch (error: any) {
            return {
                content: [
                    {
                        type: "text",
                        text: `获取进程信息时出错: ${error.message}`
                    }
                ]
            };
        }
    }
);

/**
 * 注册环境变量工具
 * 
 * 功能：获取和设置环境变量
 * 参数：
 * - action: 操作类型（get/set/list）
 * - name: 变量名（get/set操作需要）
 * - value: 变量值（set操作需要）
 * 返回值：环境变量信息
 */
server.tool(
    "environment-variables",
    "获取、设置或列出环境变量",
    {
        action: z.enum(["get", "set", "list"]).describe("操作类型"),
        name: z.string().optional().describe("变量名（get/set操作需要）"),
        value: z.string().optional().describe("变量值（set操作需要）")
    },
    async ({ action, name, value }) => {
        try {
            switch (action) {
                case "get":
                    if (!name) {
                        throw new Error("获取环境变量需要提供name参数");
                    }
                    const envValue = process.env[name];
                    return {
                        content: [
                            {
                                type: "text",
                                text: `环境变量 ${name}: ${envValue || '(未设置)'}`
                            }
                        ]
                    };

                case "set":
                    if (!name || !value) {
                        throw new Error("设置环境变量需要提供name和value参数");
                    }
                    process.env[name] = value;
                    return {
                        content: [
                            {
                                type: "text",
                                text: `环境变量已设置: ${name}=${value}`
                            }
                        ]
                    };

                case "list":
                    const envVars = Object.entries(process.env)
                        .sort(([a], [b]) => a.localeCompare(b))
                        .map(([key, val]) => `${key}=${val}`);
                    
                    return {
                        content: [
                            {
                                type: "text",
                                text: `环境变量列表:

${envVars.slice(0, 50).join('\n')}

${envVars.length > 50 ? `... 共${envVars.length}个变量，显示前50个` : `共${envVars.length}个变量`}`
                            }
                        ]
                    };
            }
        } catch (error: any) {
            return {
                content: [
                    {
                        type: "text",
                        text: `操作失败: ${error.message}`
                    }
                ]
            };
        }
    }
);

/**
 * 启动服务器的主函数
 * 
 * 该函数初始化并启动MCP服务器，通过Stdio协议进行通信
 */
async function main() {
    // 创建Stdio传输层
    const transport = new StdioServerTransport();
    
    // 连接服务器到传输层
    await server.connect(transport);
    
    // 输出启动信息到stderr（避免干扰MCP协议通信）
    console.error(`PowerShell MCP服务器已启动`);
    console.error(`操作系统: ${getPlatform()}`);
    console.error(`Node.js版本: ${process.version}`);
    console.error(`服务器版本: 1.0.0`);
    console.error(`可用的工具: execute-command, get-system-info, file-operation, get-processes, environment-variables`);
}

// 错误处理
process.on('uncaughtException', (error) => {
    console.error('未捕获的异常:', error);
});

process.on('unhandledRejection', (reason, promise) => {
    console.error('未处理的Promise拒绝:', reason);
});

// 启动服务器
main().catch((error) => {
    console.error('启动服务器失败:', error);
    process.exit(1);
});